#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"


#include "CUDA_DTIFiberTracking.h"

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>

#include <string>

#include <cutil_inline.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>

char *pDTIPathname = NULL;				/* input diffusion tensor pathname */
char *pMASKPathname = NULL;				/* input mask pathname */
char *pFiberPathname = NULL;			/* output fiber configure pathname */

float stepSize = 0.1f;
float maxLength = 50.0f;
float faThreshold = 0.1f;
float angleThreshold = 0.707f;
int scale = 1;
int innerSteps = 10;

/* input data */
CDTI *pDTI = NULL;
CScalar *pMASK = NULL;

/* MLS fiber tracking parameters */
float MLSSize = 1.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

int gpu_num = 0;
GPUPlan *plans = NULL;


/* cuda functions in cu file */
extern "C" void CUDA_BindTextures(CDTI *dti, CScalar *mask, GPUPlan *plans);
extern "C" void CUDA_GenerateSeeds(GPUPlan *plan, const float in_faThreshold, const float in_MLSSize, const int num);
extern "C" void CUDA_TrackFibers(GPUPlan *plan, const float in_stepSize, const float in_maxLength, 
								 const float in_faThreshold, const float in_angleThreshold, 
								 const float in_MLSSize, const int in_innerSteps, const int num);



void PrintUsage()
{
	printf("Usage: ");
	printf("DTIFiberTracking -dti <DTI> -mask <MASK> -fiber <Fiber> [-stepsize <stepsize> -maxlength <maxlength> -fa_thresold <fa_threshold> -angle_threshold <fa_threshold> -scale <scale>]\n");
	printf("\n");
}


int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-mask") == 0) {
			pMASKPathname = argv[++i];
		} else if (strcmp(argv[i], "-fiber") == 0) {
			pFiberPathname = argv[++i];
		} else if (strcmp(argv[i], "-stepsize") == 0) {
			stepSize = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-maxlength") == 0) {
			maxLength = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-fa_threshold") == 0) {
			faThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-angle_threshold") == 0) {
			angleThreshold = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-scale") == 0) {
			scale = atoi(argv[++i]);
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pDTIPathname == NULL || pMASKPathname == NULL || pFiberPathname == NULL)
		return -2;

	return 0;
}

void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.3f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

/* CUDA functions */
int GetMaximumGflopsDevices(int *deviceIDs)
{
	int gpu_num = 0;
	cutilSafeCall(cudaGetDeviceCount(&gpu_num));
	if(gpu_num > MAX_GPU_COUNT) gpu_num = MAX_GPU_COUNT;
	printf("\n\tCUDA-capable device count: %i\n", gpu_num);

	/* find maximum multiprocessors */
	int max_multiprocessors = 0;
	for (int deviceID = 0; deviceID < gpu_num; ++deviceID) {
		cudaDeviceProp properties;
		cutilSafeCall(cudaGetDeviceProperties(&properties, deviceID));
		if (max_multiprocessors < properties.multiProcessorCount) {
			max_multiprocessors = properties.multiProcessorCount;
		}
	}

	/* find all devices which have the maximum multiprocessors */
	int gpu_count = 0;
	for (int deviceID = 0; deviceID < gpu_num; ++deviceID) {
		cudaDeviceProp properties;
		cutilSafeCall(cudaGetDeviceProperties(&properties, deviceID));
		if (max_multiprocessors == properties.multiProcessorCount) {
			deviceIDs[gpu_count] = deviceID;
			gpu_count++;
		}
	}

	printf("\tUsing %d GPUs with %d multi processors\n\n", 
		gpu_count, max_multiprocessors);

	return gpu_count;
}

void CreateFiberTrackingPlans(const int *deviceIDs)
{
	int w = pDTI->GetWidth();
	int h = pDTI->GetHeight();
	int d = pDTI->GetDepth();
	int n = pDTI->GetN();

	// create streams, allocate memory, and copy data to device for each GPU for fiber tracking
	for(int i = 0; i < gpu_num; i++){
		plans[i].deviceID = deviceIDs[i];

		cutilSafeCall(cudaSetDevice(plans[i].deviceID));
		cutilSafeCall(cudaStreamCreate(&(plans[i].stream)));

		// allocate device memory
		cutilSafeCall(cudaMalloc((void**)&(plans[i].d_pFibers), THREADS_PER_RUN*sizeof(CUDA_Fiber)));
		cutilSafeCall(cudaMalloc((void**)&(plans[i].d_pSeeds), THREADS_PER_RUN*sizeof(CUDA_Seed)));
		
		cutilSafeCall(cudaMalloc((void**)&(plans[i].d_pSamplePoints), SAMPLE_POINT_NUM*sizeof(float3)));
		cutilSafeCall(cudaMalloc((void**)&(plans[i].d_pSampleWeights), SAMPLE_POINT_NUM*sizeof(float)));

		// copy the data from host to device
		cutilSafeCall(cudaMemcpy(plans[i].d_pSamplePoints, pSamplePoints, 
			SAMPLE_POINT_NUM*sizeof(float3), cudaMemcpyHostToDevice));
		cutilSafeCall(cudaMemcpy(plans[i].d_pSampleWeights, pSamplePointsWeight, 
			SAMPLE_POINT_NUM*sizeof(float), cudaMemcpyHostToDevice));

		// allocate host memory
		cutilSafeCall(cudaMallocHost((void**)&(plans[i].h_pFibers), THREADS_PER_RUN*sizeof(CUDA_Fiber)));
		cutilSafeCall(cudaMallocHost((void**)&(plans[i].h_pSeeds), THREADS_PER_RUN*sizeof(CUDA_Seed)));

		// create and bind textures on CUDA
		CUDA_BindTextures(pDTI, pMASK, &(plans[i]));
	}
}


void Initialize()
{
	/* creat MLS sample points */
	CreateMLSSamplePoints();
}

void Initialize_CUDA()
{
	/* get GPUs with maximum  */
	int deviceIDs[MAX_GPU_COUNT];
	gpu_num = GetMaximumGflopsDevices(deviceIDs);

	plans = new GPUPlan[gpu_num];
	memset(plans, 0, sizeof(GPUPlan)*gpu_num);
	CreateFiberTrackingPlans(deviceIDs);
}

int GenerateSeeds(CUDA_Seed *seeds, const int sw, const int sh, const float z)
{
	CUDA_Seed *tempSeeds = new CUDA_Seed[sw*sh];
	memset(tempSeeds, 0, sizeof(CUDA_Seed)*sw*sh);

	for (int y = 0; y < sh; ++y) {
		for (int x = 0; x < sw; ++x) {
			int index = y * sw + x;
			tempSeeds[index].pos.x = (float)x / (float)scale;
			tempSeeds[index].pos.y = (float)y / (float)scale;
			tempSeeds[index].pos.z = z;
			tempSeeds[index].id = index;
		}
	}

	int remain_num = sw * sh;
	int last_num = 0;
	while (remain_num > 0) {
		int this_num = THREADS_PER_RUN * gpu_num;
		if (this_num > remain_num) {
			this_num = remain_num;
		}

		/* lunch the kernel for each device */
		for(int i = 0; i < gpu_num; i++)  
		{
			memset(plans[i].h_pSeeds, 0, sizeof(CUDA_Seed)*THREADS_PER_RUN);
			int temp_num = this_num - THREADS_PER_RUN * i;
			temp_num = temp_num > THREADS_PER_RUN ? THREADS_PER_RUN : temp_num;
			if (temp_num <= 0)
				continue;
			memcpy(plans[i].h_pSeeds, &(tempSeeds[i*THREADS_PER_RUN+last_num]), sizeof(CUDA_Seed)*temp_num);
			
			CUDA_GenerateSeeds(&(plans[i]), faThreshold, MLSSize, temp_num);	
		}

		/* collect results */
		for(int i = 0; i < gpu_num; i++)  
		{
			cutilSafeCall(cudaSetDevice(plans[i].deviceID));
			cudaStreamSynchronize(plans[i].stream);
			int temp_num = this_num - THREADS_PER_RUN * i;
			temp_num = temp_num > THREADS_PER_RUN ? THREADS_PER_RUN : temp_num;
			if(temp_num <= 0)
				continue;
			memcpy(&(tempSeeds[i*THREADS_PER_RUN+last_num]), plans[i].h_pSeeds, sizeof(CUDA_Seed)*temp_num);
		}

		remain_num -= this_num;
		last_num += this_num;
	}

	int count = 0;
	for (int y = 0; y < sh; ++y) {
		for (int x = 0; x < sw; ++x) {
			int index = y * sw + x;
			if (tempSeeds[index].id >= 0) {
				seeds[count] = tempSeeds[index];
				count++;
			}
		}
	}

	SafeDeleteArray(tempSeeds);

	return count;
}

void TrackFibers(CUDA_Seed *seeds, const int count, const int sw, const int sh, const float z, CFiber *fibers)
{
	for (int y = 0; y < sh; ++y) {
		for (int x = 0; x < sw; ++x) {
			int index = y * sw + x;
			CPoint3F pos;
			pos.m_x = (float)x / (float)scale;
			pos.m_y = (float)y / (float)scale;
			pos.m_z = z;
			Seed seed;
			seed.dir = CVector3F(0, 0, 0);
			seed.pos = pos;
			fibers[index].CreateFiber(&pos, &pos, 1, 1, &seed);
		}
	}

	int remain_num = count;
	int last_num = 0;
	while (remain_num > 0) {
		int this_num = THREADS_PER_RUN * gpu_num;
		if (this_num > remain_num) {
			this_num = remain_num;
		}

		/* lunch the kernel for each device */
		for(int i = 0; i < gpu_num; i++)  
		{
			memset(plans[i].h_pSeeds, 0, sizeof(CUDA_Seed)*THREADS_PER_RUN);
			int temp_num = this_num - THREADS_PER_RUN * i;
			temp_num = temp_num > THREADS_PER_RUN ? THREADS_PER_RUN : temp_num;
			if (temp_num <= 0)
				continue;
			memcpy(plans[i].h_pSeeds, &(seeds[i*THREADS_PER_RUN+last_num]), sizeof(CUDA_Seed)*temp_num);

			CUDA_TrackFibers(&(plans[i]), stepSize, maxLength, faThreshold, angleThreshold, MLSSize, innerSteps, temp_num);
		}

		/* collect results */
		for(int i = 0; i < gpu_num; i++)  
		{
			cutilSafeCall(cudaSetDevice(plans[i].deviceID));
			cudaStreamSynchronize(plans[i].stream);
			int temp_num = this_num - THREADS_PER_RUN * i;
			temp_num = temp_num > THREADS_PER_RUN ? THREADS_PER_RUN : temp_num;
			if(temp_num <= 0)
				continue;

			for (int j = 0; j < temp_num; ++j) {
				int index = plans[i].h_pFibers[j].seed.id;
				Seed seed;
				seed.dir.m_x = plans[i].h_pFibers[j].seed.dir.x;
				seed.dir.m_y = plans[i].h_pFibers[j].seed.dir.y;
				seed.dir.m_z = plans[i].h_pFibers[j].seed.dir.z;
				seed.pos.m_x = plans[i].h_pFibers[j].seed.pos.x;
				seed.pos.m_y = plans[i].h_pFibers[j].seed.pos.y;
				seed.pos.m_z = plans[i].h_pFibers[j].seed.pos.z;
				fibers[index].CreateFiber((CPoint3F *)(plans[i].h_pFibers[j].f), 
					(CPoint3F *)(plans[i].h_pFibers[j].b), 
					plans[i].h_pFibers[j].fc, plans[i].h_pFibers[j].bc, &seed);
			}
		}

		printf("#");
		fflush(stdout);

		remain_num -= this_num;
		last_num += this_num;
	}
}

void DTIFiberTracking()
{
	int w = pDTI->GetWidth();
	int h = pDTI->GetHeight();
	int d = pDTI->GetDepth();

	int sw = w * scale;
	int sh = h * scale;
	int sd = d * scale;

	FILE *fp = fopen(pFiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, stepSize*(float)innerSteps, maxLength);

	std::string pathname = pFiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	printf("\n=============== Begin Fiber Tracking ===============\n");
	for (int z = sd/2-10; z < sd/2+10; ++z) {
		unsigned int hTimer;
		cutilCheckError(cutCreateTimer(&hTimer));
		cutilCheckError(cutResetTimer(hTimer));
		cutilCheckError(cutStartTimer(hTimer));
		printf("Slice z=%.04f ", (float)z/(float)scale);
		fflush(stdout);

		CUDA_Seed *seeds = new CUDA_Seed[sw*sh];
		int count = GenerateSeeds(seeds, sw, sh, (float)z/(float)scale);

		CFiber *fibers = new CFiber[sw*sh];
		TrackFibers(seeds, count, sw, sh, (float)z/(float)scale, fibers);

		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_z=%08.04f.conf", path.c_str(), (float)z/(float)scale);
		sprintf(vertPathname, "%sslice_z=%08.04f.vert", path.c_str(), (float)z/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);

		SaveFibers(confPathname, vertPathname, fibers, sw, sh);

		SafeDeleteArray(seeds);
		SafeDeleteArray(fibers);

		cutilCheckError(cutStopTimer(hTimer));
		printf(" Done. GPU Processing time: %f (sec)\n", cutGetTimerValue(hTimer)/1000);
	}
	printf("=============== Finish ===============\n\n");
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(pDTIPathname);
	pMASK = new CScalar();
	pMASK->OpenScalarFile(pMASKPathname);

	Initialize();

	Initialize_CUDA();

	DTIFiberTracking();

	SafeDelete(pDTI);
	SafeDelete(pMASK);

	return 0;
}